// "Copyright [2021] <Copyright Han>"

// The topology of this network is shown as the ./oxc_cluster.png

package cluster.networks;

import cluster.network_devices.*;
import cluster.task_scheduler.task_generator.*;
import cluster.network_controller.*;
import cluster.dispatcher.*;
import cluster.resource_scheduler.*;
import cluster.task_scheduler.*;
import cluster.network_controller.topo_manager.*;
import cluster.roce.*;
import inet.networklayer.configurator.ipv4.*;
import inet.node.inet.Router;

network PFCNetwork
{
    parameters:
        int numMachine = default(2);
        int numGpusperMachine = default(8);
        int switch_num = default(numMachine+numGpusperMachine);
        int numNvlink_port = default(6);
        double switch_delay = default(0);
        int inetModuleNumber = default(numGpusperMachine * numMachine);
        int TorSwitchPort = default(numGpusperMachine); //(numMachine*numGpusperMachine)/(switch_num-1)
        //int inetModuleNumber = 0;
        // enum ClusterType {
        //    OpticalCluster,
        //    ElectricalCluster,
        //    IbCluster
        //};        
        int clustertype = default(1);
        int numJobs = default(1);
        int modelsize_ = default(1000000);
        // bool useFileInput = default(false);
        int inputType = default(1);
        string file_path = default("./input_elec.txt");
        int maxnumBatches = default(1);

        int tmp_tasktype = default(0);
        
    submodules:
        AIServer[numMachine] : Machine{
            parameters:
                numGpus = numGpusperMachine;
                numNvlink_port = numNvlink_port;
                my_index = index;
        }

        topo_manager : TopoManager {}

        dispatcher : Dispatcher {}

        resource_scheduler : ResourceScheduler {
            GpuScheduleAlgo = 3;
            job_num = numJobs;
        }

        task_scheduler : TaskScheduler {
            parameters:
                maxnumGpus = numGpusperMachine * numMachine;
                modelsize = modelsize_;
                maxnumBatches = maxnumBatches;
                **.numJobs = numJobs;
                **.numGpus = int(numGpusperMachine * numMachine/numJobs);
        }

        network_controller : NetworkController {
        }

        configurator: Ipv4NetworkConfigurator {
            parameters:
                @display("p=100,100;is=s");
                dumpAddresses = true;
                dumpTopology = true;
                dumpLinks = true;
                //dumpConfig = "./networkconfig.xml";
                assignAddresses = false;
                inetModuleNum = inetModuleNumber;
                addSubnetRoutes = false;
                optimizeRoutes = false;
                networkType = 1;
        }
        adapter:RoutingAdapter {
        }
        Leaf_Switch[numMachine]: Router {
            parameters:
            enablePFC = true;
            enableECN = true;
            **.numPorts = 2*numGpusperMachine;
            **.pauseUnit = 1000000000;
        }
        Spine_Switch[numGpusperMachine]: Router {
            parameters:
            enablePFC = true;
            enableECN = true;
            **.numPorts = numMachine;
            **.pauseUnit = 1000000000;
        }

    connections allowunconnected:
       // Connect the data plane
        // Connect the machines first
        for i=0..numMachine-1, for j=0..numGpusperMachine-1 {
            AIServer[i].out_port[j] <-->Fiber<--> Leaf_Switch[i].ethg++;
        }

        for k=0..numMachine-1, for j=0..numGpusperMachine-1 {
            Leaf_Switch[k].ethg++ <-->Fiber<--> Spine_Switch[j].ethg++;
        }


        // Connect the task generator and controller
        task_scheduler.port --> dispatcher.task_generator_port;

        // Connect the controller - AIServer
        for i=0..numMachine-1 {
            dispatcher.machine_port++<-->ControlPlaneLink<-->AIServer[i].controller_port;
        }

       dispatcher.network_routing_port<--ControlPlaneLink<--network_controller.dispatcher_port$o;
 
       // Connect the switch - centralcontroller
    //    network_controller.adapter_port <-->ControlPlaneLink <--> adapter.controller_port;

       // Connect apater and switch
       for i=0..numMachine-1{
        network_controller.switch_port++ <--> ControlPlaneLink <--> Leaf_Switch[i].controlPort;
       }

       // Connect apater and switch
       for i=0..numGpusperMachine-1{
        network_controller.switch_port++ <--> ControlPlaneLink <--> Spine_Switch[i].controlPort;
       }

       // Connect apater and configurator
       configurator.networkConfiguratorOutput  -->ControlPlaneLink -->  adapter.configuration_port;
}
